package com.cupk2.llm.client.impl;

import com.cupk2.llm.client.inter.Cupk2EmbedClientInter;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.cupk2.llm.domain.LlmClient;
import org.cupk2.tools.json.CupkJsons;

import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class Cupk2EmbedOllamaClientImpl implements Cupk2EmbedClientInter {
    private LlmClient paras;

    @Override
    public String getName() {
        return "Aliyun_Embedding_SDK";
    }

    @Override
    public void init(LlmClient client) {
        paras = client;
    }

    @Override
    public float[] getEmbedding(String text) {
        HttpClient client = HttpClient.newHttpClient();
        Map<String, Object> requestBody = new HashMap<>();
        requestBody.put("model", paras.getModelkey());
        requestBody.put("input", text);


        try {
            String requestBodyString = CupkJsons.toJson(requestBody);
            HttpRequest request = HttpRequest.newBuilder()
                    .uri(URI.create(paras.getBaseurl()))
                    .header("Content-Type", "application/json")
                    .header("Authorization", "Bearer " + paras.getApikey())
                    .POST(HttpRequest.BodyPublishers.ofString(requestBodyString))
                    .build();
            HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
            if (response.statusCode() == 200) {
                float[] embeddings = getEmbeddings(response.body());
                return embeddings;
            } else {
                throw new RuntimeException("Failed to retrieve response, status code: " + response.statusCode() + ", response: " + response.body());
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    public static float[] getEmbeddings(String json) {
        ObjectMapper objectMapper = new ObjectMapper();
        try {

            Map<String, Object> responseMap = objectMapper.readValue(json, new TypeReference<Map<String, Object>>() {
            });

            @SuppressWarnings("unchecked")
            List<List<Double>> embeddings = (List<List<Double>>) responseMap.get("embeddings");

            if (embeddings == null || embeddings.isEmpty()) {
                throw new IllegalArgumentException("Embeddings field is missing or empty in the JSON response.");
            }

            List<Double> embeddingList = embeddings.get(0);

            float[] embeddingArray = new float[embeddingList.size()];
            for (int i = 0; i < embeddingList.size(); i++) {
                embeddingArray[i] = embeddingList.get(i).floatValue();
            }

            return embeddingArray;
        } catch (Exception e) {
            throw new RuntimeException("Error parsing JSON: " + e.getMessage(), e);
        }
    }
}
